from histogram_filter.particleFilter import ParticleFilter
from histogram_filter.histogramFilter import HistogramFilter
from Dataloader.src import waymo_training_plot_tools as plot_tool
from Dataloader.src import waymo_scenario_classes as class_tool
from Dataloader.src.waymo_transfer_training_to_tf import *
from Dataloader.src.waymo_training_math_tools import *
from Dataloader.src.waymo_scenario_metrics import *
from matplotlib import pyplot as plt
from driver_model.Moveable import Moveable
from driver_model.IDM import IDM
from driver_model.Velocity import VelocityModel
import tensorflow as tf
import numpy as np
import copy
import pickle

GROUND_TRUTH_STEPS = 31
TOTAL_STEPS = 91
PREDICT_STEPS = 80

idm_filter = dict()
left_lc_filter = dict()
right_lc_filter = dict()
multimodal_filter = dict()

idm_model = dict()
left_lc_model = dict()
right_lc_model = dict()
current_velocity_model = dict()

predicted_state = dict()


def filterModels(self_id):
    # construct self vehicle object
    self_state_list = dataloader.get_agent_states_list_byId(self_id)
    self_state = self_state_list[0]
    pos = [self_state[3], self_state[4]]
    vel = [self_state[10], self_state[11], self_state[12]]
    lane = 1
    length, width = self_state[6], self_state[7]
    self_vehicle = Moveable(1, pos, vel, lane, length, width)
    # self_pos_list = np.array([self_state_list[:, 3], self_state_list[:, 4]]).T
    self_pos_x = dataloader.get_agent_center_x_list_byId(self_id)
    self_pos_y = dataloader.get_agent_center_y_list_byId(self_id)
    self_pos_list = np.array([self_pos_x, self_pos_y]).T

    # construct the vehicle object to follow
    fwd_id = dataloader.get_around_cars_byID(self_id)[0][0][0]
    if fwd_id is not None:
        fwd_state_list = dataloader.get_agent_states_list_byId(fwd_id)
        fwd_state = fwd_state_list[0]
        pos = [fwd_state[3], fwd_state[4]]
        vel = [fwd_state[10], fwd_state[11], fwd_state[12]]
        lane = 1
        length, width = fwd_state[6], fwd_state[7]
        fwd = Moveable(1, pos, vel, lane, length, width)
        fwd_pos_list = np.array([fwd_state_list[:, 3], fwd_state_list[:, 4]]).T
    else:
        fwd = None

    # construct lane object
    around_lane_ids, self_lane_ids = dataloader.get_agent_all_time_lane_list_byID(self_id)
    left_lane_id = around_lane_ids[0][1][0]
    right_lane_id = around_lane_ids[0][2][0]
    self_lane_id = around_lane_ids[0][0]
    self_lane_points = dataloader.get_long_lane_points_byID(self_lane_id)
    left_lane_points = dataloader.get_long_lane_points_byID(left_lane_id)
    right_lane_points = dataloader.get_long_lane_points_byID(right_lane_id)


    delta = 1

    def idm_propagate(x, s0):
        idm_model[self_id].s0 = s0
        position_distribution, _, _, _ = idm_model[self_id].getPosDistribution(self_vehicle, fwd, self_lane_points[0], self_lane_points[1])
        predicted_position, _ = idm_model[self_id].get_position(self_vehicle, fwd, self_lane_points[0], self_lane_points[1])
        return position_distribution, predicted_position

    def left_lc_propagate(x, max_steer):
        left_lc_model[self_id].max_steer = max_steer
        predicted_position = left_lc_model[self_id].getTrajectory(self_vehicle, None, left_lane_points[0], left_lane_points[1], points_number=1)
        return None, [predicted_position[0][0], predicted_position[1][0]]

    def right_lc_propagate(x, max_steer):
        right_lc_model[self_id].max_steer = max_steer
        predicted_position = right_lc_model[self_id].getTrajectory(self_vehicle, None, right_lane_points[0], right_lane_points[1], points_number=1)
        return None, [predicted_position[0][0], predicted_position[1][0]]


    def update_vehicle_state(i):
        self_vehicle.pos = self_pos_list[i]
        self_vehicle.vel = [self_state_list[i][10], self_state_list[i][11], self_state_list[i][12]]
        if fwd is not None:
            fwd.pos = fwd_pos_list[i]
            fwd.vel = [fwd_state_list[i][10], fwd_state_list[i][11], fwd_state_list[i][12]]
            fwd.acc = (fwd_state_list[i][12] - fwd_state_list[0][12]) / (0.1 * i) * 3


    s0_particles = np.linspace(5, 15, 10)
    max_steer_particles = np.linspace(0.1, 0.3, 20)

    # construct model
    idm_model[self_id] = IDM(10, 10, 4, delta, s0=s0_particles[0], size=TOTAL_STEPS - GROUND_TRUTH_STEPS, max_steer=0.15)
    left_lc_model[self_id] = IDM(10, 10, 4, delta, s0=10, size=TOTAL_STEPS - GROUND_TRUTH_STEPS, max_steer=0.15)
    right_lc_model[self_id] = IDM(10, 10, 4, delta, s0=10, size=TOTAL_STEPS - GROUND_TRUTH_STEPS, max_steer=0.15)
    current_velocity_model[self_id] = VelocityModel()

    idm_filter[self_id] = ParticleFilter(s0_particles, idm_propagate, self_pos_list[0])
    left_lc_filter[self_id] = ParticleFilter(max_steer_particles, left_lc_propagate, self_pos_list[0])
    right_lc_filter[self_id] = ParticleFilter(max_steer_particles, right_lc_propagate, self_pos_list[0])

    model_names = ['idm', 'velocity']
    multimodal_filter[self_id] = HistogramFilter(model_names)
    for i in range(1, GROUND_TRUTH_STEPS):
        # print("timestamp: ", i)
        # print("current velocity: ", self.vel[2])
        # print("current acc: ", self.calc_acc())

        # print("Ground truth: ", self_pos_list[i])

        # filter idm parameter(s0)
        idm_filter[self_id].updateBelief(self_pos_list[i])
        left_lc_filter[self_id].updateBelief(self_pos_list[i])
        right_lc_filter[self_id].updateBelief(self_pos_list[i])
        # print(np.round(idm_filter.belief, 4))

        # update vehicle state
        update_vehicle_state(i)
        idm_model[self_id].s0 = s0_particles[np.argmax(idm_filter[self_id].belief)]
        left_lc_model[self_id].max_steer = max_steer_particles[np.argmax(left_lc_filter[self_id].belief)]
        right_lc_model[self_id].max_steer = max_steer_particles[np.argmax(right_lc_filter[self_id].belief)]

        # predict with idm
        idm_predicted_position, _  = idm_model[self_id].get_position(self_vehicle, fwd, self_lane_points[0], self_lane_points[1])
        left_lc_position = left_lc_model[self_id].getTrajectory(self_vehicle, None, left_lane_points[0], left_lane_points[1], points_number=1)
        left_lc_position = [left_lc_position[0][0], left_lc_position[1][0]]
        right_lc_position = right_lc_model[self_id].getTrajectory(self_vehicle, None, right_lane_points[0], right_lane_points[1], points_number=1)
        right_lc_position = [right_lc_position[0][0], right_lc_position[1][0]]
        # predict with velocity
        velocity_predicted_position = current_velocity_model[self_id].get_position(self_vehicle)

        # calculate the probability of each model
        multimodal_filter[self_id].updateBelief([idm_predicted_position, velocity_predicted_position, left_lc_position, right_lc_position], self_pos_list[i + 1])

        # print(multimodal_filter[self_id].histogram, ":", np.round(multimodal_filter[self_id].belief, 4))
        # print("")
        #
        # print("IDM Predicted: ", idm_predicted_position)
        # print("Velocity model predicted: ", velocity_predicted_position)

        # for full-dataset filter
    print("ID: ", self_id)
    print(multimodal_filter[self_id].belief)
    # set the model parameter to the max parameter
    idm_model[self_id].s0 = s0_particles[np.argmax(idm_filter[self_id].belief)]
    idm_trajectory = idm_model[self_id].getTrajectory(self_vehicle, fwd, self_lane_points[0], self_lane_points[1])
    left_lc_trajectory = left_lc_model[self_id].getTrajectory(self_vehicle, None, left_lane_points[0], left_lane_points[1], points_number=TOTAL_STEPS - GROUND_TRUTH_STEPS)
    right_lc_trajectory = right_lc_model[self_id].getTrajectory(self_vehicle, None, right_lane_points[0],
                                                              right_lane_points[1], points_number=TOTAL_STEPS - GROUND_TRUTH_STEPS)
    # plot_tool.visualize_on_step(dataset, dataloader, [idm_trajectory[0], idm_trajectory[1]], GROUND_TRUTH_STEPS - 1, plot_path="./")
    # plot_tool.visualize_on_step(dataset, dataloader, [right_lc_trajectory[0], right_lc_trajectory[1]], GROUND_TRUTH_STEPS - 1, plot_path="./")
    # plot_tool.visualize_on_step(dataset, dataloader, [left_lane_points[0], left_lane_points[1]], GROUND_TRUTH_STEPS - 1, plot_path="./")
    # plot_tool.visualize_on_step(dataset, dataloader, [self_lane_points[0], self_lane_points[1]], GROUND_TRUTH_STEPS - 1, plot_path="./")
    return [multimodal_filter[self_id].belief, idm_trajectory, None, left_lc_trajectory, right_lc_trajectory], self_pos_list


def calcAcc(all_agent_ids, self_id, pred_trajectory, pred_score, gt_mask, result_value, gt_value):
    score, idm_traj, current_velocity_traj, left_lc_traj, right_lc_traj = result_value
    self_index = all_agent_ids.index(self_id)
    pred_trajectory[0, self_index, 0, 0, 0:GROUND_TRUTH_STEPS - 11] = gt_value[11:GROUND_TRUTH_STEPS]
    pred_trajectory[0, self_index, 0, 0, GROUND_TRUTH_STEPS - 11:] = np.array(idm_traj).T
    pred_trajectory[0, self_index, 1, 0, 0:GROUND_TRUTH_STEPS - 11] = gt_value[11:GROUND_TRUTH_STEPS]
    pred_trajectory[0, self_index, 1, 0, GROUND_TRUTH_STEPS - 11:] = np.array(left_lc_traj).T
    pred_trajectory[0, self_index, 2, 0, 0:GROUND_TRUTH_STEPS - 11] = gt_value[11:GROUND_TRUTH_STEPS]
    pred_trajectory[0, self_index, 2, 0, GROUND_TRUTH_STEPS - 11:] = np.array(right_lc_traj).T
    score = [score[0], score[2], score[3]]
    pred_score[0, self_index] = score
    gt_mask[0, self_index, 0] = True
    return pred_trajectory, pred_score, gt_mask

if __name__ == '__main__':
    dataset = '../datasets/uncompressed_scenario_training_training.tfrecord-00166-of-01000'
    dataloader = Dataloader(dataset)

    all_agents_id_list = dataloader.get_all_agents_id()

    vehicles_to_predict = dataloader.get_agents_to_pre()



    result = dict()
    gt = dict()
    cars_to_be_predicted = []

    for car_id in cars_to_be_predicted:
        result[car_id], gt[car_id] = filterModels(car_id)

    plot_tool.visualize_on_step(dataset, dataloader, [[0], [0]], GROUND_TRUTH_STEPS - 1, plot_path="./")


    # metrics calculation
    pred_trajectory = np.zeros((1, 128, 3, 1, 80, 2), dtype=np.float32) - 1
    pred_score = np.zeros((1, 128, 3), dtype=np.float32) - 1
    gt_mask = np.zeros((1, 128, 1), dtype=bool)
    for car_id in cars_to_be_predicted:
        pred_trajectory, pred_score, gt_mask = calcAcc(all_agents_id_list, car_id, pred_trajectory, pred_score, gt_mask, result[car_id], gt[car_id])

    metric = GetMetrics(dataloader)
    metric.display_motion_metrics(pred_trajectory, pred_score, gt_mask)

    plot_tool.trajectory_visualization(dataset, dataloader, result, 1, save_path="./", start_time=GROUND_TRUTH_STEPS, is_show_gt=False)










